"""
双模态融合网络架构
结构：地震波处理分支 + 结构参数分支 → 特征融合 → 损伤预测
"""
import torch  # 添加这一行，确保torch模块被正确导入
import torch.nn as nn


class MultiModalNet(nn.Module):
    def __init__(self, num_concrete_grades=3):
        """
        参数说明：
        - num_concrete_grades: 混凝土强度类别数（默认C30/C40/C50三种）
        """
        super().__init__()

        # 地震波处理分支（1D CNN）
        self.wave_branch = nn.Sequential(
            # 输入形状：[batch, 3通道, 1000时间点]
            nn.Conv1d(3, 64, kernel_size=50, padding=25),  # 保持时序长度
            nn.BatchNorm1d(64),  # 稳定训练过程
            nn.ReLU(),
            nn.MaxPool1d(2)  # 降采样到500点
        )

        # 结构参数分支（全连接网络）
        self.struct_branch = nn.Sequential(
            # 输入特征：梁宽+梁高+配筋率 + 混凝土独热编码
            nn.Linear(3 + num_concrete_grades, 64),
            nn.ReLU(),
            nn.Linear(64, 64)
        )

        # 特征融合与回归
        self.fusion = nn.Sequential(
            nn.Linear(64 * 500 + 64, 256),  # 地震特征64*500 + 结构特征64
            nn.ReLU(),
            nn.Dropout(0.3),  # 防止过拟合
            nn.Linear(256, 1),
            nn.Sigmoid()  # 输出0-1损伤指数
        )

    def forward(self, wave, struct):
        # 处理地震波输入
        wave_feat = self.wave_branch(wave)
        wave_feat = wave_feat.view(wave.size(0), -1)  # 展平 [batch, 64*500]

        # 处理结构参数
        struct_feat = self.struct_branch(struct)

        # 特征融合
        combined = torch.cat([wave_feat, struct_feat], dim=1)
        return self.fusion(combined)